
import numpy as np
import matplotlib.pyplot as plt
# import torchvision
# from PIL import Image
# import torchvision.transforms as transforms
# import torch
# to_pil = transforms.ToPILImage()
# # # Load CIFAR-100 dataset without transforming it to tensor
# transform = transforms.Compose([transforms.ToTensor()]) 
# trainset = torchvision.datasets.CIFAR100(root='../data', train=True, download=True, transform=transform)
# Get all class names
# class_names = cifar100.classes
def filter_image(img, filter_size=2, show_filters=False):
    """Apply Fourier Transform, high-pass and low-pass filters, and return filtered images."""
    # Convert PIL image to numpy array
    img_array = np.array(img)
    
    # Apply Fourier Transform
    f_transform = np.fft.fft2(img_array)
    f_transform_shifted = np.fft.fftshift(f_transform)
    
    # Define low-pass and high-pass filters
    rows, cols = img_array.shape
    crow, ccol = rows // 2, cols // 2
    low_pass_filter = np.zeros((rows, cols))
    high_pass_filter = np.ones((rows, cols))
    
    # Create the filters
    low_pass_filter[crow-filter_size//2:crow+filter_size//2, ccol-filter_size//2:ccol+filter_size//2] = 1
    high_pass_filter[crow-filter_size//2:crow+filter_size//2, ccol-filter_size//2:ccol+filter_size//2] = 0
    
    # Apply filters
    low_passed = f_transform_shifted * low_pass_filter
    high_passed = f_transform_shifted * high_pass_filter
    
    # Inverse Fourier Transform
    low_passed_img = np.fft.ifft2(np.fft.ifftshift(low_passed)).real
    high_passed_img = np.fft.ifft2(np.fft.ifftshift(high_passed)).real
    
    # Plot filters if show_filters is True
    if show_filters:
        plt.figure(figsize=(10, 5))
        
        # Low-pass filter
        plt.subplot(1, 2, 1)
        plt.imshow(low_pass_filter, cmap='gray')
        plt.title('Low-Pass Filter')
        plt.axis('on')

        # High-pass filter
        plt.subplot(1, 2, 2)
        plt.imshow(high_pass_filter, cmap='gray')
        plt.title('High-Pass Filter')
        plt.axis('on')
        
        plt.show()
    
    return low_passed_img, high_passed_img